import os
import torch
import torch.nn as nn
import torchvision.transforms as tvt
from PIL import Image
import numpy as np
from numpy import linalg as LA
from .impl.bit import *
from .utils import to_pil


if torch.cuda.is_available():
    DEVICE = torch.device('cuda')
else:
    DEVICE = torch.device('cpu')


class BiTFeatureExtractor(object):
    FEATURE_DIM = 2048
    def __init__(self, model_name):
        # load model
        assert model_name in KNOWN_MODELS.keys()
        self.model = KNOWN_MODELS[model_name]()
        
        weights_path = os.path.join('checkpoints', '{}.npz'.format(model_name))
        print('loading weights from {}'.format(weights_path))
        self.model.load_from(np.load(weights_path))

        # modify model head.
        self.model.head = nn.AdaptiveAvgPool2d(output_size=1)
        self.model.to(DEVICE)

        # load data transformations
        crop_size = 224
        self.transforms = tvt.Compose([
            tvt.Resize((crop_size, crop_size)),
            tvt.ToTensor(),
            tvt.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ])
    
    def extract(self, image):
        image = to_pil(image)
        inputs = self.transforms(image)
        inputs = inputs.unsqueeze(0)
        with torch.no_grad():
            inputs = inputs.to(DEVICE)
            outputs = self.model(inputs)
        outputs = outputs.squeeze(0).detach().cpu().numpy()

        # TODO: do feature normalization
        outputs = outputs / LA.norm(outputs)
    
        return outputs
